import numpy as np
import pandas as pd
from sklearn.linear_model import Ridge
from sklearn.metrics import mean_squared_error
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler, PolynomialFeatures

# 1.导入数据
data_csv = pd.read_csv('data/insurance_data.csv')
# 2.处理数据：处理非数值；处理null值；选出X、Y ；分出测试集和训练集
data = pd.get_dummies(data_csv)
X = data.drop('charges', axis=1)
Y = data['charges']
X.fillna(0, inplace=True)
Y.fillna(0, inplace=True)
X_train, X_test, Y_train, Y_test = train_test_split(X, Y, test_size=0.3)
# 3.归一化处理
scaler = StandardScaler()
X_train_scaler = scaler.fit_transform(X_train)
X_test_scaler = scaler.transform(X_test)
# 4.多项式升维
features = PolynomialFeatures(degree=2)
X_test_scaler = features.fit_transform(X_test_scaler)
X_train_scaler = features.fit_transform(X_train_scaler)
# 5.用L2做正则化求解
rge = Ridge(alpha=10)
rge.fit(X_train_scaler, np.log1p(Y_train))

y_pred_train = rge.predict(X_train_scaler)
y_pred_test = rge.predict(X_test_scaler)
# 7.评估
train_mse = mean_squared_error(y_true=np.log1p(Y_train), y_pred=y_pred_train)
test_mse = mean_squared_error(y_true=np.log1p(Y_test), y_pred=y_pred_test)
print('训练集评估log1p：', train_mse)
print('测试集评估log1p：', test_mse)
train_mse = np.sqrt(mean_squared_error(y_true=np.log1p(Y_train), y_pred=y_pred_train))
test_mse = np.sqrt(mean_squared_error(y_true=np.log1p(Y_test), y_pred=y_pred_test))
print('训练集评估sqrt(log1p)：', train_mse)
print('测试集评估sqrt(log1p)：', test_mse)
train_mse = mean_squared_error(y_true=Y_train, y_pred=y_pred_train)
test_mse = mean_squared_error(y_true=Y_test, y_pred=y_pred_test)
print('训练集评估：', train_mse)
print('测试集评估：', test_mse)
